import numpy as np
import acado
import time
import matplotlib.pyplot as plt
from numpy import sin, cos,exp

import time
import rospy

N=10 #timestep
NU=4
NOD=3
NX=16
NY=16
NYN=4

NW=16
NWN=12

x0=np.zeros(NX).reshape(1,NX)
od=np.zeros(NOD)
y=np.zeros(NY)
u=np.zeros(NU)
Q=np.eye(NW)

Q=np.diag([1e-8,1e-8,1e-8,1,1,1,1,100,10,10,100,100,1,5,5,5e5])


W=np.transpose(np.tile(Q,N))

print(W.shape)

WN=Q[:-4,:-4]

y=np.zeros([N,NY])


x0[0,0]=-2.3 #px
x0[0,1]=0.02#py
x0[0,2]=0.67 #pz
x0[0,3]=1.0#qw
x0[0,4]=0.0#qx
x0[0,5]=0.0#qy
x0[0,6]=0.0#qz
x0[0,7]=0.0#vx
x0[0,8]=0.0#vy
x0[0,9]=0.0#vz
x0[0,10]=1.0#pxx
x0[0,11]=0.0#pxy
x0[0,12]=0.0#pxz
x0[0,13]=1.0#pyy
x0[0,14]=0.0#pyz
x0[0,15]=1.0#pzz



X=np.tile(x0.reshape(NX,1),N+1).T

print(X.shape)

OD=np.ones([N+1,NOD])
for i in range(N):
    OD[i,0:3]=np.array([3.3,-1.3,0.67])+np.array([0.01,0,0])*i




    

u=np.zeros([N,NU])

xlog=[]
tracelog=[]

t0=time.time()


for i in range(5000):
    
    OD[:,0]=3.3+i*0.01
    OD[:,1]=sin(i*0.01)*0.0
    #print("iter: "+str(i))
    X,u=acado.mpc(0,1,x0,X,u,OD,y,W,WN,1)
    x0=X[1,:].reshape(1,NX)
    P=np.array([[x0[0,10],x0[0,11],x0[0,12]],
                [x0[0,11],x0[0,13],x0[0,14]],
                [x0[0,12],x0[0,14],x0[0,15]]])
    
    
    #tracelog.append(x0[0,10]+x0[0,13]+x0[0,15])
    tracelog.append(np.linalg.det(P))
    
    xlog.append(x0)
    xlog_arr=np.vstack(xlog)
    
    plt.subplot(211)
    plt.plot(xlog_arr[:,0],xlog_arr[:,1],"-r")
    plt.gcf().canvas.mpl_connect('key_release_event',
                                             lambda event: [exit(
                                                 0) if event.key == 'escape' else None])

    plt.scatter(OD[0,0],OD[0,1])
    
    plt.subplot(212)
    plt.title('traceP')
    plt.plot(tracelog)
    plt.ylim([-0.01,0.2])
    
    if i%10==0:
        plt.pause(0.0001)

t1=time.time()

print(t1-t0)


plt.subplot(221)
plt.plot(xlog_arr[:,0])
plt.subplot(222)
plt.plot(xlog_arr[:,1])
plt.subplot(223)
plt.plot(tracelog)

plt.title('tracking result')



    
 
    
plt.grid()
plt.subplot(224)
plt.plot(xlog_arr[:,0],xlog_arr[:,1])
plt.title('determinat of covariance')

plt.show()

